# encoding: utf-8
import sys
from io import TextIOWrapper
import os
import re
from functional import seq


class InstructionType:
    A_INSTRUCTION = 1
    C_INSTRUCTION = 2
    L_INSTRUCTION = 3


class Constructor():
    def __init__(self, rfile: TextIOWrapper):
        self.rfile = rfile
        self.rfile_size = os.path.getsize(self.rfile.name)
        self.instruction = ""

    def hasMoreLines(self) -> bool:
        return self.rfile.tell() != self.rfile_size

    def advance(self) -> str:
        if self.hasMoreLines():
            line = self.rfile.readline().strip("\n")
            # comment line
            if re.match(r'//.*', line):
                return self.advance()
            # blank line
            elif re.match(r'^\s*$', line):
                return self.advance()

            if "//" in line:
                line = line[0:line.index("//")]
            return ''.join(line.split(' '))
        else:
            return ""

    def instructionType(self):
        if '@' in self.instruction:
            return InstructionType.A_INSTRUCTION
        elif '(' in self.instruction:
            return InstructionType.L_INSTRUCTION
        else:
            return InstructionType.C_INSTRUCTION

    def symbol(self):
        # When commandType() is A_COMMAND or L_COMMAND
        # return the xxx in @xxx or (xxx)
        if self.instructionType() == InstructionType.A_INSTRUCTION:
            return self.instruction[1:]
        elif self.instructionType() == InstructionType.L_INSTRUCTION:
            return self.instruction[1:-1]
        else:
            return ""

    # C code 三种情况
    # dest = comp; jump
    # comp; jump
    # dest = comp
    def ccode(self, line: str):
        codes = []
        if ';' not in line:
            codes.extend(line.split('='))
            codes.append("")    # jump == ""
        elif '=' not in line:
            codes.append("")    # dest == ""
            codes.extend(line.split(';'))
        else:
            cons = line.split('=')
            codes.append(cons[0])
            codes.extend(cons[1].split(';'))

        return codes

    def dest(self):
        if self.instructionType() == InstructionType.C_INSTRUCTION:
            return self.ccode(self.instruction)[0]
        else:
            return ""

    def comp(self):
        if self.instructionType() == InstructionType.C_INSTRUCTION:
            return self.ccode(self.instruction)[1]
        else:
            return ""

    def jump(self):
        if self.instructionType() == InstructionType.C_INSTRUCTION:
            return self.ccode(self.instruction)[2]
        else:
            return ""


class Code:
    def a_instr(self, code: int) -> str:
        bins = bin(code)[2:]
        return '0' * (16 - len(bins)) + bins

    def dest(self, code: str) -> str:
        bin = ['0', '0', '0']
        if 'A' in code:
            bin[0] = '1'
        if 'D' in code:
            bin[1] = '1'
        if 'M' in code:
            bin[2] = '1'

        return ''.join(bin)

    def comp(self, code: str) -> str:
        comps_dict = {'0': '101010', '1': '111111', '-1': '111010', 'D': '001100', 'A': '110000', '!D': '001101', '!A': '110001', '-D': '001111', '-A': '110011', 'D+1': '011111', 'A+1': '110111', 'D-1': '001110', 'A-1': '110010', 'D+A': '000010',
                      'D-A': '010011', 'A-D': '000111', 'D&A': '000000', 'D|A': '010101', 'M': '110000', '!M': '110001', '-M': '110011', 'M+1': '110111', 'M-1': '110010', 'D+M': '000010', 'D-M': '010011', 'M-D': '000111', 'D&M': '000000', 'D|M': '010101'}

        if code not in comps_dict.keys():
            # 可交换顺序的操作符
            if len(code) == 3:
                op = code[1]
                if op in ['&', '|', '+']:
                    code = code[::-1]
        return comps_dict[code]


    def jump(self, code: str) -> str:
        jump_dict = {
            '':    '000',
            'JGT': '001',
            'JEQ': '010',
            'JGE': '011',
            'JLT': '100',
            'JNE': '101',
            'JLE': '110',
            'JMP': '111',
        }

        return jump_dict[code]


class SymTable:
    def __init__(self):
        self.symTable: dict[str, int] = {
            "R0": 0,
            "R1": 1,
            "R2": 2,
            "R3": 3,
            "R4": 4,
            "R5": 5,
            "R6": 6,
            "R7": 7,
            "R8": 8,
            "R9": 9,
            "R10": 10,
            "R11": 11,
            "R12": 12,
            "R13": 13,
            "R14": 14,
            "R15": 15,
            "SP": 0,
            "LCL": 1,
            "ARG": 2,
            "THIS": 3,
            "THAT": 4,
            "SCREEN": 16384,
            "KBD": 24576,
        }

    def addEntry(self, symbol: str, address: int):
        self.symTable[symbol] = address

    def contains(self, symbol: str) -> bool:
        return True if symbol in self.symTable else False

    def getAddress(self, symbol: str) -> int:
        return self.symTable[symbol]

def translate(filename, hackname):
    print("translate {}".format(filename))
    codec = Code()
    symTable = SymTable()

    # The first parse: add symbol of L_INSTRUCTION
    with open(filename, 'r') as rfile:
        cons = Constructor(rfile)
        address = 0
        while cons.hasMoreLines():
            cons.instruction = cons.advance()
            if cons.instructionType() == InstructionType.L_INSTRUCTION:
                symTable.addEntry(cons.symbol(), address)
                address -= 1
            address += 1

    # The second parse: encode
    with open(filename, 'r') as rfile, open(hackname, 'w') as wfile:
        cons = Constructor(rfile)
        code = ""
        symbol_index = 16
        while cons.hasMoreLines():
            cons.instruction = cons.advance()
            if cons.instruction == "":
                break
            if cons.instructionType() == InstructionType.A_INSTRUCTION:
                sym = cons.symbol()
                if sym.isdigit():
                    wfile.write(codec.a_instr(int(sym)) + "\n")
                else:
                    if not symTable.contains(sym):
                        symTable.addEntry(sym, symbol_index)
                        wfile.write(codec.a_instr(int(symbol_index)) + "\n")
                        symbol_index += 1
                    else:
                        wfile.write(codec.a_instr(
                            symTable.getAddress(sym)) + "\n")
            elif cons.instructionType() == InstructionType.C_INSTRUCTION:
                comp = cons.comp()
                dest = cons.dest()
                jump = cons.jump()
                comp_code = codec.comp(comp)
                dest_code = codec.dest(dest)
                jump_code = codec.jump(jump)
                a = '0'

                if 'M' in comp:
                    a = '1'

                wfile.write('111' + a + comp_code +
                            dest_code + jump_code + "\n")

def main():
    files = sys.argv[1:]
    seq(files).map(lambda file: translate(file, file.rstrip('\.asm') + '.hack2')).to_list()

if __name__ == '__main__':
    main()

